package com.zhongfu.ai.web;

import cn.hutool.core.map.MapUtil;
import cn.hutool.json.JSONUtil;
import com.google.common.collect.Lists;
import com.zhongfu.ai.common.response.Result;
import com.zhongfu.ai.vo.ChatResponseVo;
import com.zhongfu.ai.vo.MetadataFileVo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RestController
@CrossOrigin("*")
@RequestMapping("/api/ollama")
public class ChatController {

    @Autowired
    private OllamaChatModel ollamaChatModel;

    @Autowired
    private OpenAiChatModel openAiChatModel;

    @Autowired
    private PgVectorStore pgVectorStore;

    @RequestMapping("/openai/chat")
    public Result openAiChat(@RequestParam String message) {
        String result = openAiChatModel.call(new Prompt(new UserMessage(message))).getResult().getOutput().getText();
        return Result.success(result);
    }

    @RequestMapping("/chat")
    public String chat(@RequestParam String message) {
        String prompt = "你是一名小助手，请使用中文回答问题：" + message;
        ChatResponse chatResponse = ollamaChatModel.call(new Prompt(new UserMessage(message)));

        return chatResponse.getResult().getOutput().getText();
    }

    @RequestMapping("/chat_stream")
    public Flux<ChatResponse> chatStream(@RequestParam String message) {
        String prompt = "你是一名小助手，请使用中文回答问题：" + message;
        return ollamaChatModel.stream(new Prompt(new UserMessage(message)));
    }


    @RequestMapping("/rag/search")
    public String searchRag(@RequestParam String ragTag,
                            @RequestParam String message) {
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;

        // 指定搜索文档
//        SearchRequest searchRequest =
//                SearchRequest.builder()
//                        .filterExpression("knowledge == '" + ragTag + "'")
//                        .topK(5)
//                        .build();

        SearchRequest searchRequest = SearchRequest.builder().query(message).build();

        List<Document> documentList = pgVectorStore.similaritySearch(searchRequest);

        String documentCollector = documentList.stream().map(Document::getText).collect(Collectors.joining());

        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentCollector));

        List<Message> messageList = new ArrayList<>();
        messageList.add(new UserMessage(message));
        messageList.add(ragMessage);

        return openAiChatModel.call(new Prompt(messageList)).getResult().getOutput().getText();
    }


    @RequestMapping("/generate_stream_rag")
    public Flux<ChatResponseVo> generateStreamRag(@RequestParam(required = false) String ragTag,
                                                @RequestParam String message) {
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;

        // 指定搜索文档
        SearchRequest searchRequest = SearchRequest.builder().query(message).build();

        List<Document> documentList = pgVectorStore.similaritySearch(searchRequest);

        documentList.stream().forEach(document ->{
            log.info("document metadata : {}",document.getMetadata());
        });
        String documentCollector = documentList.stream().map(Document::getText).collect(Collectors.joining());

        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentCollector));

        List<Message> messageList = new ArrayList<>();
        messageList.add(new UserMessage(message));
        messageList.add(ragMessage);

        Flux<ChatResponse> chatResponseFlux =  openAiChatModel.stream(new Prompt(ragMessage, new UserMessage(message)));

        Flux<ChatResponseVo> responseVoFlux = chatResponseFlux.map(chatResponse -> {
            ChatResponseVo chatResponseVo= new ChatResponseVo();
            if(chatResponse.getResult().getMetadata().getFinishReason().equalsIgnoreCase("STOP")){
              //  chatResponseVo.setFinished(true);
            }
            chatResponseVo.setType("TEXT");
            chatResponseVo.setContent(chatResponse.getResult().getOutput().getText());
            return chatResponseVo;
        });

        ChatResponseVo fileResponseVo = new ChatResponseVo();
        fileResponseVo.setType("FILE");
        fileResponseVo.setFinished(true);

        List<MetadataFileVo> files = new ArrayList<>();
        files.add(new MetadataFileVo("二年级英语辅导资料.pdf","http://xxx.sss.com/te/test.pdf","pdf"));
        files.add(new MetadataFileVo("关于开展春季数学模拟竞猜的通知.doc","http://xxx.sss.com/te/test2.doc","doc"));
        files.add(new MetadataFileVo("初中物理三年级下学期教学课件.ppt","http://xxx.sss.com/te/test3.ppt","ppt"));
        files.add(new MetadataFileVo("XXX中学七年级上学期期末考试成绩汇总.xls","http://xxx.sss.com/te/test3.xls","xls"));

        fileResponseVo.setContent(files);

        Flux<ChatResponseVo> finalChatResponseVo = responseVoFlux.concatWith(Flux.just(fileResponseVo));
        log.info(JSONUtil.toJsonStr(finalChatResponseVo));

        return finalChatResponseVo;
//        List<MetadataFileVo> files = new ArrayList<>();
//        files.add(new MetadataFileVo("测试数据.pdf","http://xxx.sss.com/te/test.pdf"));
//        files.add(new MetadataFileVo("测试数据2.pdf","http://xxx.sss.com/te/test2.pdf"));
//        files.add(new MetadataFileVo("测试数据3.pdf","http://xxx.sss.com/te/test3.pdf"));
//
//        Generation generation = new Generation(AssistantMessage.builder().content("测试内容").build());
//
//        ChatResponse fileChatResponse = ChatResponse.builder()
//                .generations(Lists.newArrayList(generation))
//                .build();
//
//        Flux<ChatResponse> newFluxChatResponse = chatResponseFlux.concatWith(Flux.just(fileChatResponse));
//        return newFluxChatResponse;
    }
}
